// Package middleware 提供请求幂等中间件，通过 Redis 记录首次响应并在重复请求时直接返回
package middleware

import (
    "bytes"
    "crypto/rand"
    "crypto/sha256"
    "encoding/hex"
    "io"
    "net/http"
    "time"

    "github.com/we-pay-kit/we-pay-kit/internal/cache"
)

// IdemConfig 配置幂等中间件
type IdemConfig struct {
    TTL time.Duration // 结果缓存 TTL
}

// IdempotencyMiddleware 幂等中间件实现
type IdempotencyMiddleware struct {
    cfg    IdemConfig
    redis  *cache.RedisClient
}

// NewIdempotency 创建幂等中间件
func NewIdempotency(cfg IdemConfig) *IdempotencyMiddleware {
    return &IdempotencyMiddleware{cfg: cfg, redis: cache.GetRedis()}
}

// onlyForMutations 仅对非 GET/HEAD 请求生效
func (m *IdempotencyMiddleware) onlyForMutations(method string) bool {
    return method != http.MethodGet && method != http.MethodHead
}

// capturedResponse 捕获的响应内容
type capturedResponse struct {
    Status int               `json:"status"`
    Header map[string]string `json:"header"`
    Body   []byte            `json:"body"`
}

// responseRecorder 简单的响应记录器
type responseRecorder struct {
    http.ResponseWriter
    statusWritten bool
    statusCode    int
    buf           bytes.Buffer
}

func (rr *responseRecorder) WriteHeader(statusCode int) {
    if !rr.statusWritten {
        rr.statusWritten = true
        rr.statusCode = statusCode
        rr.ResponseWriter.WriteHeader(statusCode)
    }
}

func (rr *responseRecorder) Write(p []byte) (int, error) {
    if !rr.statusWritten {
        rr.WriteHeader(http.StatusOK)
    }
    rr.buf.Write(p)
    return rr.ResponseWriter.Write(p)
}

// Middleware 幂等中间件：
// - 从头 `Idempotency-Key` 读取幂等键，若无则生成并回写响应头
// - 利用 Redis.SetNX 作为“进行中”信号，再将首个响应缓存 TTL 时间
// - 重复请求优先返回已缓存响应；如仅存在进行中信号则返回 409 提示等待
func (m *IdempotencyMiddleware) Middleware(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        if m.redis == nil || !m.onlyForMutations(r.Method) {
            next.ServeHTTP(w, r)
            return
        }

        idemKey := r.Header.Get("Idempotency-Key")
        if idemKey == "" {
            // 基于请求体与路径构造一个稳定 key，并在响应头返回
            bodyBytes, _ := io.ReadAll(r.Body)
            _ = r.Body.Close()
            r.Body = io.NopCloser(bytes.NewReader(bodyBytes))
            sum := sha256.Sum256(append([]byte(r.Method+"|"+r.URL.Path+"|"), bodyBytes...))
            idemKey = hex.EncodeToString(sum[:])
            // 同时返回一个随机ID，便于客户端复用
            var rnd [16]byte
            _, _ = rand.Read(rnd[:])
            w.Header().Set("Idempotency-Generated", hex.EncodeToString(rnd[:]))
        }
        w.Header().Set("Idempotency-Key", idemKey)
        keyResult := "idem:result:" + idemKey
        keyLock := "idem:lock:" + idemKey

        // 若已有结果，直接返回
        var cached capturedResponse
        if err := m.redis.GetJSON(keyResult, &cached); err == nil && cached.Status != 0 {
            for k, v := range cached.Header {
                w.Header().Set(k, v)
            }
            w.WriteHeader(cached.Status)
            _, _ = w.Write(cached.Body)
            return
        }

        // 尝试锁定进行中
        if ok, _ := m.redis.SetNX(keyLock, "1", m.cfg.TTL); !ok {
            // 已有请求在进行，返回 409 提示重试
            w.WriteHeader(http.StatusConflict)
            _, _ = w.Write([]byte("request in progress, please retry later"))
            return
        }

        // 包装响应，处理完成后写入缓存
        rr := &responseRecorder{ResponseWriter: w}
        next.ServeHTTP(rr, r)

        // 采集必要头（仅复制常用内容类型）
        headers := map[string]string{}
        if ct := w.Header().Get("Content-Type"); ct != "" {
            headers["Content-Type"] = ct
        } else {
            headers["Content-Type"] = "application/json"
        }
        res := capturedResponse{Status: rr.statusCode, Header: headers, Body: rr.buf.Bytes()}
        _ = m.redis.SetJSON(keyResult, res, m.cfg.TTL)
        // 锁键保持原 TTL，无需显式删除
    })
}


